{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.font_manager as mfm\n",
    "import os\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "from datetime import timedelta \n",
    "import yfinance as yf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "_CHN_FONT_ = None\n",
    "_FONT_PROP_ = None\n",
    "_CHN_EN_DICT_ = './data/locationDict.csv'\n",
    "_CHN_LatLng_DICT_ = './data/lat_lng.csv'\n",
    "# The province names in China\n",
    "# Ignore Hong Kong, Macau, Taiwan and Tibet\n",
    "# English name\n",
    "names_province = ['Anhui', 'Beijing', 'Chongqing', 'Fujian', 'Gansu', 'Guangdong', 'Guangxi', 'Guizhou',\n",
    "                  'Hainan', 'Hebei', 'Heilongjiang', 'Henan', 'Hubei', 'Hunan', 'Inner Mongolia', 'Jiangsu',\n",
    "                  'Jiangxi', 'Jilin', 'Liaoning', 'Ningxia', 'Qinghai', 'Shaanxi', 'Shandong', 'Shanghai',\n",
    "                  'Shanxi', 'Sichuan', 'Tianjin', 'Xinjiang', 'Yunnan', 'Zhejiang']\n",
    "# Chinese name\n",
    "names_province_cn = ['安徽省', '北京市', '重庆市', '福建省', '甘肃省', '广东省', '广西壮族自治区', '贵州省',\n",
    "                     '海南省', '河北省', '黑龙江省', '河南省', '湖北省', '湖南省', '内蒙古自治区', '江苏省',\n",
    "                     '江西省', '吉林省', '辽宁省', '宁夏回族自治区', '青海省', '陕西省', '山东省', '上海市',\n",
    "                     '山西省', '四川省', '天津市', '新疆维吾尔自治区', '云南省', '浙江省']\n",
    "\n",
    "# A one to one corrspondence with the names used by Pyecharts\n",
    "pyecharts_province_dict = {'北京市':'北京', '天津市':'天津', '河北省':'河北', '山西省':'山西', \n",
    "                           '内蒙古自治区':'内蒙古', '辽宁省':'辽宁', '吉林省':'吉林', '黑龙江省':'黑龙江', \n",
    "                           '上海市':'上海', '江苏省':'江苏', '浙江省':'浙江', '安徽省':'安徽', \n",
    "                           '福建省':'福建', '江西省':'江西', '山东省':'山东', '河南省':'河南', \n",
    "                           '湖北省':'湖北','湖南省':'湖南', '广东省':'广东', '广西壮族自治区':'广西', \n",
    "                           '海南省':'海南', '重庆市':'重庆', '四川省':'四川', '贵州省':'贵州', \n",
    "                           '云南省':'云南', '西藏自治区':'西藏', '陕西省':'陕西', '甘肃省':'甘肃', \n",
    "                           '青海省':'青海', '宁夏回族自治区':'宁夏', '新疆维吾尔自治区':'新疆', \n",
    "                           '香港特别行政区':'香港', '澳门特别行政区':'澳门','台湾省':'台湾'}\n",
    "\n",
    "# English names of the capital cities of provinces\n",
    "provincial_capital_dict = {'Anhui':'Hefei', 'Fujian':'Fuzhou', 'Gansu':'Lanzhou','Guangdong':'Guangzhou','Guizhou':'Guiyang', \n",
    "                           'Hainan':'Haikou', 'Hebei':'Shijiazhuang','Heilongjiang':'Harbin', 'Henan':'Zhengzhou',\n",
    "                           'Hubei':'Wuhan', 'Hunan':'Changsha', 'Jiangsu':'Nanjing', 'Jiangxi':'Nanchang', \n",
    "                           'Jilin':'Changchun', 'Liaoning':'Shenyang', 'Qinghai':'Xining', 'Shaanxi':'Xi’an', \n",
    "                           'Shandong':'Jinan', 'Shanxi':'Taiyuan', 'Sichuan':'Chengdu', 'Yunnan':'Kunming', \n",
    "                           'Zhejiang':'Hangzhou', 'Guangxi':'Nanning','Inner Mongolia':'Hohhot', 'Ningxia':'Yinchuan',\n",
    "                           'Xinjiang':'Ürümqi', 'Tibet':'Lhasa', 'Beijing':'Beijing', 'Chongqing':'Chongqing', \n",
    "                           'Shanghai':'Shanghai', 'Tianjin':'Tianjin'}\n",
    "\n",
    "# Populations of provinces\n",
    "# data source: http://www.chamiji.com/2019chinaprovincepopulation\n",
    "# unit 10,000\n",
    "provincial_population_dict = {'Anhui':6323.6, 'Fujian':3941, 'Gansu': 2637.26,'Guangdong':11346,'Guizhou':3600, \n",
    "                           'Hainan':934.32, 'Hebei':7556.30,'Heilongjiang':3773.1, 'Henan':9605,\n",
    "                           'Hubei':5917, 'Hunan':6898.8, 'Jiangsu':8050.7, 'Jiangxi':4647.6, \n",
    "                           'Jilin':2704.06, 'Liaoning':4359.3, 'Qinghai':603.23, 'Shaanxi':3864.4, \n",
    "                           'Shandong':10047.24, 'Shanxi':3718.34, 'Sichuan':8341, 'Yunnan':4800.5, \n",
    "                           'Zhejiang':5737, 'Guangxi':4926,'Inner Mongolia':2534.0, 'Ningxia':688.11,\n",
    "                           'Xinjiang':2486.76, 'Tibet':337.15, 'Beijing':2154.2, 'Chongqing': 3101.79, \n",
    "                           'Shanghai':2423.78, 'Tianjin':1559.60}\n",
    "# Populations of cities\n",
    "# data source: http://www.chamiji.com/\n",
    "# unit 10,000\n",
    "city_population_dict = {'北京市':2154.2, '天津市':1559.60, '上海市':2423.78, '重庆市':3101.79,\n",
    "                        '亳州市':523.7, '六安市':483.7, '合肥市':808.7, '安庆市':469.1, '宣城市':264.8, '宿州市':568.14, # 安徽省\n",
    "                        '池州市':147.4, '淮北市':225.4, '淮南市':349.0, '滁州市':411.4, '芜湖市':374.8, '蚌埠市':339.20,\n",
    "                        '铜陵市':162.9, '阜阳市':820.7, '马鞍山市':233.7, '黄山市':140.7,\n",
    "                        '南京市':843.62, '南通市':731, '宿迁市':492.59, '常州市':472.9, '徐州市':880.20, '扬州市':453.1, # 江苏省\n",
    "                        '无锡市':657.45, '泰州市':463.57, '淮安市':492.50, '盐城市':720, '苏州市':1072.17, '连云港市':452.0,\n",
    "                        '镇江市':319.64,\n",
    "                        '东营市':217.21, '临沂市':1062.4, '威海市':283, '德州市':581, '日照市':293.03, '枣庄市':392.73, # 山东省\n",
    "                        '泰安市':564.0, '济南市':746.04, '济宁市':834.59, '淄博市':470.2, '滨州市':392.25, '潍坊市':937.3,\n",
    "                        '烟台市':712.18, '聊城市':606.43, '菏泽市':876.5, '青岛市':939.48, \n",
    "                        '丽水市':219.25, '台州市':613.90, '嘉兴市':472.6, '宁波市':820.2, '杭州市':980.6, '温州市':925, # 浙江省\n",
    "                        '湖州市':301.09, '绍兴市':503.5, '舟山市':117.3, '衢州市':220.9, '金华市':560.4, \n",
    "                        '三明市':258, '南平市':269, '厦门市':411, '宁德市':291, '泉州市':870, '漳州市':514.0, # 福建省\n",
    "                        '福州市':774, '莆田市':290.0, '龙岩市':264, \n",
    "                        '东莞市':834.25, '中山市':326, '云浮市':326, '佛山市':765.67, '广州市':1449.84, '惠州市':477.70, # 广东省\n",
    "                        '揭阳市':608.6, '梅州市':437.43, '汕头市':560.82, '汕尾市':305.33, '江门市':456.17, '河源市':309.11,\n",
    "                        '深圳市':1252.83, '清远市':386.0, '湛江市':730.5, '潮州市':265.08, '珠海市':176.54, '肇庆市':411.54, \n",
    "                        '茂名市':620.41, '阳江市': 254.29, '韶关市':297.92, \n",
    "                        '北海市':166.33, '南宁市':715.33, '崇左市':208.68, '来宾市':221.86, '柳州市':400.00, '桂林市':505.75, # 广西壮族自治区\n",
    "                        '梧州市':303.7, '河池市':352.35, '玉林市': 581.08, '百色市':364.65, '贵港市':437.54, '贺州市':205.67,\n",
    "                        '钦州市':328, '防城港市':94.02, \n",
    "                        '万宁市':float(\"NaN\"), '三亚市':76.42, '儋州市':90.57, '海口市':227.21, '三沙市':float(\"NaN\"), '东方市':float(\"NaN\"), # 海南省 ########## \n",
    "                        '临高县':float(\"NaN\"), '乐东黎族自治县':float(\"NaN\"),'五指山市':float(\"NaN\"),'保亭黎族苗族自治县':float(\"NaN\"), '定安县':float(\"NaN\"),'屯昌县':float(\"NaN\"),\n",
    "                        '文昌市':float(\"NaN\"),'昌江黎族自治县':float(\"NaN\"),'澄迈县':float(\"NaN\"),'琼中黎族苗族自治县':float(\"NaN\"),'琼海市':float(\"NaN\"),'白沙黎族自治县':float(\"NaN\"),\n",
    "                        '陵水黎族自治县':float(\"NaN\"),\n",
    "                        '保定市':1046.92, '唐山市': 789.7, '廊坊市':467.8, '张家口市':443.3, '承德市':356.50, '沧州市':755.49, # 河北省\n",
    "                        '石家庄市':1087.99, '秦皇岛市':311.08, '衡水市':446.04, '邢台市':735.16, '邯郸市':951.11, \n",
    "                        '乌兰察布市':210.25, '乌海市':56.11, '兴安盟':160.42, '包头市':287.8, '呼伦贝尔市':252.92, '呼和浩特市':311.5, # 内蒙古自治区\n",
    "                        '巴彦淖尔市':168.5, '赤峰市':431.5, '通辽市':312.87, '鄂尔多斯市':206.87, '锡林郭勒盟':105.16, '阿拉善盟':24.8,\n",
    "                        '临汾市':450.03, '吕梁市':388.56, '大同市':345.60, '太原市':442.15, '忻州市':317.20, '晋中市':338.15, # 山西省\n",
    "                        '晋城市':234.31, '朔州市':178.12, '运城市':535.97, '长治市':346.8, '阳泉市':141.44, \n",
    "                        '三门峡市':227.29, '信阳市':647.41, '南阳市':1001.36, '周口市':867.78, '商丘市':732.53, '安阳市':517.6, # 河南省\n",
    "                        '平顶山市':520.77, '开封市':456.49, '新乡市':579.41, '洛阳市':688.85, '济源市':73.27, '漯河市':266.53,\n",
    "                        '濮阳市':360.94, '焦作市':359.07, '许昌市':443.74, '郑州市':1013.6, '驻马店市':619.02, '鹤壁市':162.73,\n",
    "                        '娄底市':393.18, '岳阳市':579.71, '常德市':582.7, '张家界市':153.79, '怀化市':497.96, '株洲市':402.08, # 湖南省\n",
    "                        '永州市':545.21, '湘潭市':286.5, '湘西土家族苗族自治州':264.95, '益阳市':441.38, '衡阳市':724.34, '邵阳市':737.05,\n",
    "                        '郴州市':474.5, '长沙市':815.47, \n",
    "                        '仙桃市':154.45, '十堰市':341.8, '咸宁市':253.51, '天门市':128.35, '孝感市':491.50, '宜昌市':413.56, # 湖北省 ##########\n",
    "                        '恩施土家族苗族自治州':336.10, '武汉市':1089.29, '潜江市': float(\"NaN\"), '神农架林区':7.68, '荆州市':564.17, '荆门市':290.15,\n",
    "                        '襄阳市':565.4, '鄂州市':107.69, '随州市':221.05, '黄冈市':634.1, '黄石市':247.05, \n",
    "                        '上饶市':678.34, '九江市':487.33, '南昌市':546.35, '吉安市':494.19, '宜春市':555.37, '抚州市':403.10, # 江西省 2018\n",
    "                        '新余市':118.07, '景德镇市':166.49, '萍乡市':192.50, '赣州市':861.2, '鹰潭市':116.75, \n",
    "                        '六盘水市':292.41, '安顺市':234.44, '毕节市':665.97, '贵阳市':480.20, '遵义市':624.83, '铜仁市':315.69, # 贵州省 2018\n",
    "                        '黔东南苗族侗族自治州':352.37, '黔南布依族苗族自治州':327.1, '黔西南布依族苗族自治州':286,\n",
    "                        '乐山市':327.21, '内江市':375.37, '凉山彝族自治州':521.29, '南充市':641.79, '宜宾市':453, '巴中市':331.67, # 四川省 2018\n",
    "                        '广元市':266.00, '广安市':325.0, '德阳市':353.2, '成都市':1604.5, '攀枝花市':123.61, '泸州市':431.72,\n",
    "                        '甘孜藏族自治州':118.63, '眉山市':297.48, '绵阳市':483.56, '自贡市':290.14, '资阳市':255.3, '达州市':568.95,\n",
    "                        '遂宁市':323.59, '阿坝藏族羌族自治州':94.01, '雅安市':153.78, \n",
    "                        '临沧市':252.60, '丽江市':129.0, '保山市':261.4, '大理白族自治州':361.88, '德宏傣族景颇族自治州':130.90, '怒江傈僳族自治州':54.7, # 云南省 2018\n",
    "                        '文山壮族苗族自治州':363.6, '昆明市':678.3, '昭通市': 553.7, '普洱市':262.7, '曲靖市':612.2, '楚雄彝族自治州':274.40,\n",
    "                        '玉溪市':238.1, '红河哈尼族彝族自治州':471.3, '西双版纳傣族自治州':118.0, '迪庆藏族自治州':41.2,\n",
    "                        '咸阳市':437.6, '商洛市':238.13, '安康市':266.1, '宝鸡市':378.10, '延安市':226.31, '榆林市':340.33, # 陕西省\n",
    "                        '汉中市':344.93, '渭南市':538.29, '西安市':953.44, '铜川市':83.34,\n",
    "                        '临夏回族自治州':204.41, '兰州市':372.96, '嘉峪关市':24.98, '天水市':333.98, '定西市':280.84, '平凉市':211.28, # 甘肃省\n",
    "                        '庆阳市':200.55, '张掖市':122.93, '武威市':182.53, '甘南藏族自治州':74.23, '白银市':172.93, '酒泉市':112.36,\n",
    "                        '金昌市':46.92, '陇南市':287.42, \n",
    "                        '果洛藏族自治州':20.57, '海东市':147.08, '海北藏族自治州':28.3, '海南藏族自治州':47.24, '海西蒙古族藏族自治州':51.52, '玉树藏族自治州':40.95, # 青海省\n",
    "                        '西宁市':235.50, '黄南藏族自治州':27.42, \n",
    "                        '乌鲁木齐市':222.61, '五家渠市':float(\"NaN\"), '伊犁哈萨克自治州':461.71, '克孜勒苏柯尔克孜自治州':62.06, '克拉玛依市':44.28, '北屯市':float(\"NaN\"), # 新疆维吾尔自治区 2018 ##########\n",
    "                        '博尔塔拉蒙古自治州':47.54, '双河市':float(\"NaN\"), '可克达拉市':float(\"NaN\"), '吐鲁番市':63.73, '和田地区':252.28, '哈密市':56.11,\n",
    "                        '喀什地区':464.97, '图木舒克市':float(\"NaN\"), '塔城地区':15.2, '巴音郭楞蒙古自治州':127.93, '昆玉市':float(\"NaN\"), '昌吉回族自治州':161,\n",
    "                        '石河子市':float(\"NaN\"), '铁门关市':float(\"NaN\"), '阿克苏地区':254.6, '阿勒泰地区':67.16, '阿拉尔市':float(\"NaN\"),\n",
    "                        '中卫市':115.75, '吴忠市':140.37, '固原市':122.82, '石嘴山市':80.29, '银川市':222.54, # 宁夏回族自治区 2018\n",
    "                        '吉林市':415.35, '四平市':320.4, '延边朝鲜族自治州':210.14, '松原市':275.41, '白城市':190.9, '白山市':119.5, # 吉林省 2018\n",
    "                        '辽源市':117.94, '通化市':217.15, '长春市':748.9,\n",
    "                        '丹东市':239.5, '大连市':698.75, '抚顺市':210.7, '朝阳市':295, '本溪市':147.63, '沈阳市':829.4, # 辽宁省 2018\n",
    "                        '盘锦市':143.65, '营口市':243.8, '葫芦岛市':277.0, '辽阳市':183.7, '铁岭市':299.8, '锦州市':296.4,\n",
    "                        '阜新市':186.2, '鞍山市':344.0,\n",
    "                        '七台河市':78.6, '伊春市':115.9, '佳木斯市':234.5, '双鸭山市':142.3, '哈尔滨市':955.0, '大兴安岭地区':43.93, # 黑龙江省 2018\n",
    "                        '大庆市':273.1, '牡丹江市':254.8, '绥化市':527.6, '鸡西市':175, '鹤岗市':100.9, '黑河市':160.5,\n",
    "                        '齐齐哈尔市':533.7\n",
    "                       }\n",
    "\n",
    "# Reset the unit of populations to be 1 instead of 10,000\n",
    "provincial_population_dict.update((x, y*1e4) for x, y in provincial_population_dict.items())\n",
    "city_population_dict.update((x, y*1e4) for x, y in city_population_dict.items())\n",
    "\n",
    "# Obtain the latitude and longitude of a province from the local dictionary\n",
    "# Save as a dictionary\n",
    "def get_province_latlng_dict():\n",
    "    province_latlng_dict = {}\n",
    "    chn_latlng = pd.read_csv(_CHN_LatLng_DICT_, encoding='utf-8')\n",
    "    for key in provincial_capital_dict:\n",
    "        lat = chn_latlng[chn_latlng.city == provincial_capital_dict[key]]['lat'].tolist()[0]\n",
    "        lng = chn_latlng[chn_latlng.city == provincial_capital_dict[key]]['lng'].tolist()[0]\n",
    "        province_latlng_dict[key] = (lat, lng)\n",
    "    return province_latlng_dict\n",
    "\n",
    "# Save as a dataframe\n",
    "def get_province_latlng_df():\n",
    "    data = pd.DataFrame.from_dict(province_latlng_dict, orient='index', columns = ['lat', 'lng'])\n",
    "    data = data.reset_index()\n",
    "    data = data.rename(columns={'index':'province_name'})\n",
    "    #data_gephi = data.rename(columns={'province_name':'id'})\n",
    "    #data_gephi.to_csv(r'./data/data_network_P2P_gephi_id.csv', index = False)\n",
    "    return data\n",
    "\n",
    "province_latlng_dict = get_province_latlng_dict()\n",
    "province_latlng_df = get_province_latlng_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings on the Chinese font\n",
    "def set_font(font_file):\n",
    "    if not os.path.exists(font_file):\n",
    "        print(font_file + \" not found.  If you wish to display Chinese characters in plots, please use set_font() to set the path to the font file.\")\n",
    "    else:\n",
    "        global _CHN_FONT_, _FONT_PROP_\n",
    "        _CHN_FONT_ = font_file\n",
    "        _FONT_PROP_ = mfm.FontProperties(fname=_CHN_FONT_)\n",
    "    return\n",
    "\n",
    "set_font('./STFANGSO.TTF')   # for displaying Chinese characters in plots\n",
    "\n",
    "def use_chn():\n",
    "    return _CHN_FONT_ is None\n",
    "\n",
    "# Add English name for a province or a city\n",
    "def add_en_location(df, tag = 'city'):\n",
    "    '''Add province_name_en, and city_name_en'''\n",
    "    chn_en = pd.read_csv(_CHN_EN_DICT_, encoding='utf-8')\n",
    "    translation = dict([t for t in zip(chn_en['Chinese'], chn_en['English'])])\n",
    "    if tag == 'province':\n",
    "        df['province_name_en'] = df['province_name'].replace(translation)\n",
    "    elif tag == 'city':\n",
    "        df['province_name_en'] = df['province_name'].replace(translation)\n",
    "        df['city_name_en'] = df['city_name'].replace(translation)\n",
    "    elif tag == 'network':\n",
    "        df['source_en'] = df['source'].replace(translation)\n",
    "        df['target_en'] = df['target'].replace(translation)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data loading\n",
    "_Data_PATH_ = './data/parameters/'\n",
    "_Data_PATH_X_ = './data/parameters_x/'\n",
    "_Data_PATH_XX_ = './data/parameters_xx/'\n",
    "_Data_PATH_M_ = './data/parameters_m/'\n",
    "_Data_PATH_MX_ = './data/parameters_mx/'\n",
    "_Data_PATH_MXX_ = './data/parameters_mxx/'\n",
    "_City_PATH_ = './data/data_DXY_city_all.csv'\n",
    "_Province_PATH_ = './data/data_DXY_province_all.csv'\n",
    "_Province_Domestic_PATH_ = './data/data_DXY_province_all_domestic.csv'\n",
    "_Network_P2P_PATH_ = './data/data_network_P2P.csv'\n",
    "\n",
    "def load_DXY_raw():\n",
    "    raw_city = pd.read_csv(_City_PATH_)\n",
    "    raw_province = pd.read_csv(_Province_PATH_) \n",
    "    raw_province_domestic = pd.read_csv(_Province_Domestic_PATH_) \n",
    "    raw_city['update_date'] = pd.to_datetime(raw_city['update_date'])\n",
    "    raw_city['update_date'] = raw_city['update_date'].dt.date \n",
    "    raw_province['update_date'] = pd.to_datetime(raw_province['update_date'])\n",
    "    raw_province['update_date'] = raw_province['update_date'].dt.date\n",
    "    raw_province_domestic['update_date'] = pd.to_datetime(raw_province_domestic['update_date'])\n",
    "    raw_province_domestic['update_date'] = raw_province_domestic['update_date'].dt.date\n",
    "    return raw_city, raw_province, raw_province_domestic\n",
    "\n",
    "def load_network_raw():\n",
    "    # Do not distinguish between move_in and move_out\n",
    "    raw = pd.read_csv(_Network_P2P_PATH_) \n",
    "    raw['update_date'] = pd.to_datetime(raw['update_date'])\n",
    "    raw['update_date'] = raw['update_date'].dt.date \n",
    "    return raw\n",
    "\n",
    "def load_ind_simulation_raw():\n",
    "    df_parameters_list = []\n",
    "    df_estimation_list = []\n",
    "    for name in names_province:\n",
    "        df_parameters_single = pd.read_csv(_Data_PATH_ + name + '_parameters.csv')\n",
    "        df_estimation_single = pd.read_csv(_Data_PATH_ + name + '_estimation.csv')\n",
    "        df_parameters_list.append(df_parameters_single)\n",
    "        df_estimation_list.append(df_estimation_single)\n",
    "    return df_parameters_list, df_estimation_list \n",
    "\n",
    "\n",
    "def load_ind_simulation_raw_x():\n",
    "    df_parameters_list = []\n",
    "    df_estimation_list = []\n",
    "    for name in names_province:\n",
    "        df_parameters_single = pd.read_csv(_Data_PATH_X_ + name + '_parameters.csv')\n",
    "        df_estimation_single = pd.read_csv(_Data_PATH_X_ + name + '_estimation.csv')\n",
    "        df_parameters_list.append(df_parameters_single)\n",
    "        df_estimation_list.append(df_estimation_single)\n",
    "    return df_parameters_list, df_estimation_list \n",
    "\n",
    "def load_ind_simulation_raw_xx():\n",
    "    df_parameters_list = []\n",
    "    df_estimation_list = []\n",
    "    for name in names_province:\n",
    "        df_parameters_single = pd.read_csv(_Data_PATH_XX_ + name + '_parameters.csv')\n",
    "        df_estimation_single = pd.read_csv(_Data_PATH_XX_ + name + '_estimation.csv')\n",
    "        df_parameters_list.append(df_parameters_single)\n",
    "        df_estimation_list.append(df_estimation_single)\n",
    "    return df_parameters_list, df_estimation_list  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some utility functions\n",
    "# Get the list of dates from the start_date to the end_date\n",
    "def daterange(start_date, end_date):\n",
    "    for n in range(int ((end_date - start_date).days) + 1):\n",
    "        yield start_date + timedelta(n)\n",
    "        \n",
    "# Transform the migration edgelist to the migration matrix\n",
    "def matrix_P2P(data_single, m):\n",
    "    mindex = np.zeros((m, m))\n",
    "    for i, name_i in enumerate(names_province):\n",
    "        for j, name_j in enumerate(names_province):\n",
    "            if i == j:\n",
    "                mindex[i][j] = 0\n",
    "            else:\n",
    "                temp = data_single[(data_single.source_en == name_i) & (data_single.target_en == name_j)].value.tolist()\n",
    "                if temp == []:\n",
    "                    mindex[i][j] = 0\n",
    "                else:\n",
    "                    mindex[i][j] = temp[0]\n",
    "    return mindex\n",
    "\n",
    "# Perform the operation for all migration data from the start_date to the end_date\n",
    "def matrix_P2P_all(data_network_P2P, start_date, end_date):\n",
    "    mindex_list = []\n",
    "    m = len(names_province)\n",
    "    for item in daterange(start_date, end_date):\n",
    "        data_single = data_network_P2P[data_network_P2P.update_date == item]\n",
    "        mindex = matrix_P2P(data_single, m)\n",
    "        mindex_list.append(mindex) \n",
    "    return mindex_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The state names in the US (including D.C.)\n",
    "# Ignore oversea territories: American Samoa, Guam, Northern Mariana Islands, Puerto Rico, and Virgin Islands\n",
    "\n",
    "names_state = ['Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut', 'Delaware', \n",
    "                 'District of Columbia',\n",
    "                 'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois', 'Indiana', 'Iowa', 'Kansas', \n",
    "                 'Kentucky', 'Louisiana', 'Maine', 'Maryland', 'Massachusetts', 'Michigan', 'Minnesota', 'Mississippi',\n",
    "                 'Missouri', 'Montana', 'Nebraska', 'Nevada', 'New Hampshire', 'New Jersey', 'New Mexico', 'New York',\n",
    "                 'North Carolina', 'North Dakota', 'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island', 'South Carolina',\n",
    "                 'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington', 'West Virginia',\n",
    "                 'Wisconsin', 'Wyoming']\n",
    "\n",
    "# The abbreviations of state names\n",
    "names_state_short_dict = {\n",
    "    'Alabama': 'AL', 'Alaska':'AK', 'Arizona':'AZ', 'Arkansas':'AR', 'California':'CA', 'Colorado':'CO', 'Connecticut':'CT', 'Delaware':'DE', \n",
    "    'Florida':'FL', 'Georgia':'GA', 'Hawaii':'HI', 'Idaho':'ID', 'Illinois':'IL', 'Indiana':'IN', 'Iowa':'IA', 'Kansas':'KS', \n",
    "    'Kentucky':'KY', 'Louisiana':'LA', 'Maine':'ME', 'Maryland':'MD', 'Massachusetts':'MA', 'Michigan':'MI', 'Minnesota':'MN', 'Mississippi':'MS',\n",
    "    'Missouri':'MO', 'Montana':'MT', 'Nebraska':'NE', 'Nevada':'NV', 'New Hampshire':'NH', 'New Jersey':'NJ', 'New Mexico':'NM', 'New York':'NY',\n",
    "    'North Carolina':'NC', 'North Dakota':'ND', 'Ohio':'OH', 'Oklahoma':'OK', 'Oregon':'OR', 'Pennsylvania':'PA', 'Rhode Island':'RI', 'South Carolina':'SC',\n",
    "    'South Dakota':'SD', 'Tennessee':'TN', 'Texas':'TX', 'Utah':'UT', 'Vermont':'VT', 'Virginia':'VA', 'Washington':'WA', 'West Virginia':'WV',\n",
    "    'Wisconsin':'WI', 'Wyoming':'WY', \n",
    "    'District of Columbia': 'DC', 'Guam': 'Guam', 'Northern Mariana Islands':'Northern Mariana Islands',\n",
    "    'Puerto Rico': 'PR', 'Virgin Islands': 'VI'}\n",
    "\n",
    "# The populations of states\n",
    "state_population_dict = {\n",
    "    'Alabama': 4903185, 'Alaska':731545, 'Arizona':7278717, 'Arkansas':3017825, \n",
    "    'California':39512223, 'Colorado':5758736, 'Connecticut':3565287, 'Delaware':973764, \n",
    "    'Florida':21477737, 'Georgia':10617423, 'Hawaii':1415872, 'Idaho':1787065, \n",
    "    'Illinois':12671821, 'Indiana':6732219, 'Iowa':3155070, 'Kansas':2913314, \n",
    "    'Kentucky':4467673, 'Louisiana':4648794, 'Maine':1344212, 'Maryland':6045680, \n",
    "    'Massachusetts':6949503, 'Michigan':9986857, 'Minnesota':5639632, 'Mississippi':2976149,\n",
    "    'Missouri':6137428, 'Montana':1068778, 'Nebraska':1934408, 'Nevada':3080156, \n",
    "    'New Hampshire':1359711, 'New Jersey':8882190, 'New Mexico':2096829, 'New York':19453561,\n",
    "    'North Carolina':10488084, 'North Dakota':762062, 'Ohio':11689100, 'Oklahoma':3956971, \n",
    "    'Oregon':4217737, 'Pennsylvania':12801989, 'Rhode Island':1059361, 'South Carolina':5148714,\n",
    "    'South Dakota':884659, 'Tennessee':6833174, 'Texas':28995881, 'Utah':3205958, \n",
    "    'Vermont':623989, 'Virginia':8535519, 'Washington':7614893, 'West Virginia':1792147,\n",
    "    'Wisconsin':5822434, 'Wyoming':578759, 'District of Columbia': 705749, 'Guam': 165718, \n",
    "    'Puerto Rico': 3193694, 'Virgin Islands': 104914}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The PVI values of states\n",
    "df_state_pvi = pd.DataFrame(columns = ['state_name', 'partisan', 'value'])\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Alabama', 'R', 14]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Alaska', 'R', 9]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Arizona', 'R', 5]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Arkansas', 'R', 15]\n",
    "\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['California', 'D', 12]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Colorado', 'D', 1]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Connecticut', 'D', 6]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Delaware', 'D', 6]\n",
    "\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Florida', 'R', 2]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Georgia', 'R', 5]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Hawaii', 'D', 18]\n",
    "\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Idaho', 'R', 19]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Illinois', 'D', 7]\n",
    "\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Indiana', 'R', 9]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Iowa', 'R', 3]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Kansas', 'R', 13]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Kentucky', 'R', 15]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Louisiana', 'R', 11]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Maine', 'D', 3]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Maryland', 'D', 12]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Massachusetts', 'D', 12]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Michigan', 'D', 1]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Minnesota', 'D', 1]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Mississippi', 'R', 9]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Missouri', 'R', 9]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Montana', 'R', 11]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Nebraska', 'R', 14]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Nevada', 'D', 1]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['New Hampshire', 'Even', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['New Jersey', 'D', 7]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['New Mexico', 'D', 3]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['New York', 'D', 12]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['North Carolina', 'R', 3]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['North Dakota', 'R', 17]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Ohio', 'R', 3]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Oklahoma', 'R', 20]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Oregon', 'D', 5]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Pennsylvania', 'Even', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Rhode Island', 'D', 10]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['South Carolina', 'R', 8]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['South Dakota', 'R', 14]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Tennessee', 'R', 14]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Texas', 'R', 8]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Utah', 'R', 20]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Vermont', 'D', 15]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Virginia', 'D', 1]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Washington', 'D', 7]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['West Virginia', 'R', 19]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Wisconsin', 'Even', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Wyoming', 'R', 25]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['District of Columbia', 'D', 43]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Guam', 'None', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Northern Mariana Islands', 'None', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Puerto Rico', 'None', 0]\n",
    "df_state_pvi.loc[len(df_state_pvi)] = ['Virgin Islands', 'None', 0]\n",
    "\n",
    "def color_pvi(df_state_pvi, state):\n",
    "    \n",
    "    palette_D = plt.get_cmap('Blues')\n",
    "    palette_R = plt.get_cmap('Reds')\n",
    "    palette_Even = plt.get_cmap('Greens')\n",
    "    palette_None = plt.get_cmap('Purples')\n",
    "    \n",
    "    D_max = max(df_state_pvi[df_state_pvi['partisan'] == 'D']['value'])\n",
    "    D_min = min(df_state_pvi[df_state_pvi['partisan'] == 'D']['value'])\n",
    "    R_max = max(df_state_pvi[df_state_pvi['partisan'] == 'R']['value'])\n",
    "    R_min = min(df_state_pvi[df_state_pvi['partisan'] == 'R']['value'])\n",
    "    \n",
    "    partisan = df_state_pvi[df_state_pvi['state_name'] == state]['partisan'].tolist()[0]\n",
    "    value = df_state_pvi[df_state_pvi['state_name'] == state]['value'].tolist()[0]\n",
    "    if partisan == 'Even':\n",
    "        return palette_Even(0.6)\n",
    "    elif partisan == 'None':\n",
    "        return palette_None(0.6)\n",
    "    elif partisan == 'D':\n",
    "        return palette_D((value - D_min)/2/(D_max - D_min) + 0.3)\n",
    "    else:\n",
    "        return palette_R((value - R_min)/2/(R_max - R_min) + 0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stock market data\n",
    "# S & P 500 index\n",
    "def load_SP_raw():\n",
    "    '''\n",
    "    This provides a way to lookinto the 'raw' data\n",
    "    '''\n",
    "    SP = yf.Ticker(\"^GSPC\")\n",
    "    raw = SP.history(period=\"max\")\n",
    "    raw = raw.reset_index()\n",
    "    rename_dict = {'Date': 'update_date',\n",
    "                   'Open': 'open',\n",
    "                   'High': 'high',\n",
    "                   'Low': 'low',\n",
    "                   'Close': 'close',\n",
    "                   'Volume': 'volume',\n",
    "                   'Dividends': 'dividends',\n",
    "                   'Stock Splits': 'stock splits'\n",
    "                  }\n",
    "    data = raw.rename(columns=rename_dict)\n",
    "    data['update_date'] = pd.to_datetime(data['update_date'])  # original type of update_time after read_csv is 'str'\n",
    "    data['update_date'] = data['update_date'].dt.date\n",
    "    data = data.reset_index(drop=True)\n",
    "    print('Data date range: ', data['update_date'].min(), 'to', data['update_date'].max())\n",
    "    print('Number of rows in raw data: ', data.shape[0])\n",
    "    return data\n",
    "# Stock price of a certain cooperation\n",
    "def load_STOCK_raw(name):\n",
    "    '''\n",
    "    This provides a way to lookinto the 'raw' data\n",
    "    '''\n",
    "    SP = yf.Ticker(name)\n",
    "    raw = SP.history(period=\"max\")\n",
    "    raw = raw.reset_index()\n",
    "    rename_dict = {'Date': 'update_date',\n",
    "                   'Open': 'open',\n",
    "                   'High': 'high',\n",
    "                   'Low': 'low',\n",
    "                   'Close': 'close',\n",
    "                   'Volume': 'volume',\n",
    "                   'Dividends': 'dividends',\n",
    "                   'Stock Splits': 'stock splits'\n",
    "                  }\n",
    "    data = raw.rename(columns=rename_dict)\n",
    "    data['update_date'] = pd.to_datetime(data['update_date'])  # original data type of update_time is 'str'\n",
    "    data['update_date'] = data['update_date'].dt.date\n",
    "    data = data.reset_index(drop=True)\n",
    "    #print('Data date range: ', data['update_date'].min(), 'to', data['update_date'].max())\n",
    "    #print('Number of rows in raw data: ', data.shape[0])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ODE systems\n",
    "#%%latex\n",
    "#Difference equations \n",
    "#without immigration\n",
    "#\\begin{align}\n",
    "#S(t + 1) - S(t) &= -\\beta(t)S(t)\\frac{I(t)}{N(t)}, \\\\\n",
    "#E(t + 1) - E(t) &= \\beta(t)S(t)\\frac{I(t)}{N(t)} - \\sigma(t)E(t), \\\\\n",
    "#I(t + 1) - I(t) &= \\sigma(t)E(t) - \\gamma(t)I(t), \\\\\n",
    "#R(t + 1) - R(t) &= \\gamma(t)I(t)\n",
    "#\\end{align}\n",
    "#Difference equations \n",
    "#with immigration\n",
    "#\\begin{align}\n",
    "#S_i(t + 1) - S_i(t) &= -\\beta_i(t)S_i(t)\\frac{I_i(t)}{N_i(t)} - \\sum_{j, j \\neq i}a_{ij}(t)S_i(t) + \\sum_{j, j \\neq i}a_{ji}(t)S_j(t) \\\\\n",
    "#E_i(t + 1) - E_i(t) &= \\beta_i(t)S_i(t)\\frac{I_i(t)}{N_i(t)} - \\sigma_i(t)E_i(t) - \\sum_{j, j \\neq i}a_{ij}(t)E_i(t) + \\sum_{j, j \\neq i}a_{ji}(t)E_j(t) \\\\\n",
    "#I_i(t + 1) - I_i(t) &= \\sigma_i(t)E_i(t) - \\gamma_i(t)I_i(t) \\\\\n",
    "#R_i(t + 1) - R_i(t) &= \\gamma_i(t)I_i(t)\n",
    "#\\end{align}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
